# -*- coding: utf-8 -*-

import os
import numpy as np
import cv2
import h5py


def show_ground_truth(ground_truth, name="color_map"):
    min_value, max_value = np.min(ground_truth), np.max(ground_truth)
    img = ((ground_truth - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)
    img = cv2.resize(img, (img.shape[1]*4, img.shape[0]*4))
    color_map = cv2.applyColorMap(img, cv2.COLORMAP_JET)
    cv2.imshow(name, color_map)


def read_labeled_points(label_file_path):
    with open(label_file_path) as label_file:
        str_ground_truths = label_file.readlines()
    labeled_points = [tuple([int(e) for e in line.strip().split(" ")]) for line in str_ground_truths]
    return labeled_points


def show_labeled_points(image_path, label_file_path):
    image = cv2.imread(image_path)
    image = image[40:, :, :]
    labeled_points = read_labeled_points(label_file_path)
    for point in labeled_points:
        if point[1] <= 40:
            continue
        cv2.circle(image, (point[0], point[1]-40), 3, (0, 0, 255), -1)
    cv2.imshow("show", image)
    cv2.waitKey()
    cv2.destroyAllWindows()


def generate_mask():
    image_mask = np.zeros(shape=(520, 960, 3), dtype=np.uint8)
    image_mask[...] = 255
    cv2.line(image_mask, (325, 0), (0, 475), (0, 0, 0), 2)
    cv2.line(image_mask, (500, 0), (960, 250), (0, 0, 0), 2)
    cv2.floodFill(image_mask, None, (160, 230), (0, 0, 0))
    cv2.floodFill(image_mask, None, (850, 130), (0, 0, 0))
    ground_truth_mask = cv2.resize(image_mask, dsize=(240, 130))
    ground_truth_mask = ground_truth_mask[:, :, 0].astype(dtype=np.float32) / 255
    return image_mask, ground_truth_mask


def generate_ground_truth(labeled_file_path, image_path, ground_truth_mask, show=False):
    image = cv2.imread(image_path)
    image = image[40:, :, :]
    image = cv2.resize(image, dsize=(image.shape[1]//2, image.shape[0]//2))
    shape = image.shape[0:2]
    shape = [shape[0]//4, shape[1]//4]
    labeled_points = read_labeled_points(labeled_file_path)
    ground_truth = np.zeros(shape=shape, dtype=np.float32)
    for location in labeled_points:
        if location[1] <= 40:
            continue
        temp = np.zeros_like(ground_truth)
        temp[(location[1]-40)//8, location[0]//8] = 1
        temp = cv2.GaussianBlur(temp, (0, 0), 5)
        temp = np.multiply(temp, ground_truth_mask)
        ground_truth += temp
    ground_truth *= 1000
    if show:
        show_ground_truth(ground_truth)
        cv2.waitKey()
        cv2.destroyAllWindows()
    return ground_truth


def main():
    src_folder = "/home/lijun/Dataset/metro/2016_09_28_17_19"
    src_images_name = [name for name in os.listdir(src_folder) if "jpg" in name]
    src_label_files_name = [name.replace("jpg", "txt") for name in src_images_name]
    image_mask, ground_truth_mask = generate_mask()
    save_folder = "/home/lijun/Dataset/metro/train/2016_09_28_17_19"
    save_index = 1
    for image_name, label_file_name in zip(src_images_name, src_label_files_name):
        image_path = os.path.join(src_folder, image_name)
        label_file_path = os.path.join(src_folder, label_file_name)
        # show_labeled_points(image_path, label_file_path)
        image = cv2.imread(image_path)[40:, :, :]
        image = cv2.resize(image, dsize=(image.shape[1]//2, image.shape[0]//2))
        assert image.shape == image_mask.shape
        image &= image_mask
        # cv2.imshow("image", src_image)
        ground_truth = generate_ground_truth(label_file_path, image_path, ground_truth_mask)
        # save crop image and ground truth
        img_height, img_width = image.shape[0], image.shape[1]
        ground_truth_height = ground_truth.shape[0]
        ground_truth_width = ground_truth.shape[1]
        assert ground_truth_height == 130 and ground_truth_width == 240
        assert img_height == 520 and img_width == 960
        img_crop_height, img_crop_width = img_height // 2, img_width // 2
        ground_truth_crop_height, ground_truth_crop_width = ground_truth_height // 2, ground_truth_width // 2
        for y in range(0, 2):
            img_offset_y = y * (img_height // 2)
            ground_truth_offset_y = y * (ground_truth_height // 2)
            for x in range(0, 2):
                img_offset_x = x * (img_width // 2)
                ground_truth_offset_x = x * (ground_truth_width // 2)
                img_crop = image[img_offset_y: img_offset_y + img_crop_height,
                                 img_offset_x: img_offset_x + img_crop_width, :]
                ground_truth_crop = ground_truth[
                                    ground_truth_offset_y: ground_truth_offset_y + ground_truth_crop_height,
                                    ground_truth_offset_x: ground_truth_offset_x + ground_truth_crop_width]
                img_save_path = "{}/images/{:>07}.jpg".format(save_folder, save_index)
                ground_truth_save_path = "{}/ground_truth/{:>07}.h5".format(save_folder, save_index)
                cv2.imwrite(img_save_path, img_crop)
                h5_file = h5py.File(ground_truth_save_path, "w")
                h5_file.create_dataset("data", data=ground_truth_crop)
                h5_file.close()
                save_index += 1
                img_save_path = "{}/images/{:>07}.jpg".format(save_folder, save_index)
                ground_truth_save_path = "{}/ground_truth/{:>07}.h5".format(save_folder, save_index)
                img_crop = cv2.flip(img_crop, flipCode=1)  # horizontal flip
                ground_truth_crop = cv2.flip(ground_truth_crop, flipCode=1)  # horizontal flip
                cv2.imwrite(img_save_path, img_crop)
                h5_file = h5py.File(ground_truth_save_path, "w")
                h5_file.create_dataset("data", data=ground_truth_crop)
                h5_file.close()
                save_index += 1
        print(save_index // 8)


def main2():
    import random
    folder = "/home/lijun/Dataset/metro/train/2016_09_27_07_09/images"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    for _ in range(1, 20):
        name = random.choice(images_name)
        image_path = os.path.join(folder, name)
        mat_path = image_path.replace("images", "ground_truth").replace("jpg", "h5")
        image = cv2.imread(image_path)
        ground_truth_file = h5py.File(mat_path, "r")
        ground_truth = np.array(ground_truth_file["data"])
        print(np.sum(ground_truth))
        ground_truth_file.close()
        cv2.imshow("image", image)
        print(np.min(ground_truth), np.max(ground_truth))
        show_ground_truth(ground_truth)
        cv2.waitKey()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
